#!/usr/bin/env python

import struct

UNPACK_MAP = {1:'B', 2:'H', 4:'L', 8:'Q'}
ENDIAN_MAP = {'big':'>', 'little':'<'}

class BitsIO:
    '''Deals with bytes objects in an IO style

    indicies are in units of bits'''

    def __init__(self,byte_str=b''):
        self.ba = bytearray(byte_str)
        self.pos = 0

    def read(self,length,**kwargs):
        start = self.pos
        self.pos = end = start+length
        if self.pos > self.bitlength():
            return(None)
        bcut = self[start:end]
        return(bcut.unpack1(**kwargs))

    def readlist(self,fmtlist):
        data = []
        for fmt in fmtlist:
            if isinstance(fmt,dict):
                sfmt = dict((str(k),v) for k,v in fmt.items())
                d = self.read(fmt['bitlength'],endian='little',**sfmt)
            elif isinstance(fmt,int):
                d = self.read(fmt,signed=False,endian='big')
            elif isinstance(fmt,list):
                # TODO how best to make this read multiple times?
                sublist = []
                d = self.readlist(fmt)
                while any(d):
                    sublist.append(d)
                    d = self.readlist(fmt)
                d = sublist
            data.append(d)
        return(data)

    def seek(self,pos):
        self.pos = pos

    def __getitem__(self,sl):
        byte_start = sl.start//8
        if (sl.stop > self.bitlength()):
            stop = self.bitlength()
        else:
            stop = sl.stop
        byte_end = (stop-1)//8 + 1

        length = stop - sl.start 
        shiftnum = 7 - (stop-1)%8 
        
        tmpba = self.ba[byte_start:byte_end]
        tmpba = shift(tmpba,shiftnum)
        tmpba = cut(tmpba,length)
        # for i,b in enumerate(self.bstr[-byte_end:byte_start]):
        #for i,b in enumerate(self.bstr[byte_start:byte_end]):
            # print(hex(ord(b)))
            # print(length, shift)

        return(BitsIO(tmpba))

    def extend(self,other):
        self.ba.extend(other.ba)

    def unpack1(self, signed=False, endian='big',**others):
        fmt = UNPACK_MAP[self.bytelength()]
        if signed:
            fmt = fmt.lower()
        fmt = ENDIAN_MAP[endian] + fmt
        return(struct.unpack(fmt,str(self))[0])
        
    def bytelength(self):
        return(len(self.ba))
    
    def bitlength(self):
        return(self.bytelength()*8)

    def __hex__(self):
        return("0x"+("".join('{0:02x}'.format(x)
                             for x in self.ba).lstrip('0')
                     or '0'))
    def __bytes__(self):
        return bytes(self.ba)

    def __str__(self):
        return str(self.ba)

def cut(ba, length):
    byte_length = (length-1)//8 + 1
    tmpba = ba[-byte_length:]
    last_byte_length = length%8 or 8
    tmpba[0] &= ones(last_byte_length)
    return(tmpba)

def shift(ba,num):
    out = bytearray()
    ba = bytearray(ba)
    carry = 0
    for b in ba:
        out += chr((b+carry) >> num)
        carry = (b & ones(num)) << 8
    return(out)

def ones(length):
    return((1<<length)-1)
